package com.pethospital.config;

import com.fasterxml.jackson.databind.ObjectMapper;
import com.pethospital.dto.MessageDTO;
import com.pethospital.entity.Message;
import com.pethospital.service.ConsultService;
import com.pethospital.util.JwtUtil;
import io.jsonwebtoken.Claims;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.web.socket.*;
import org.springframework.web.socket.config.annotation.EnableWebSocket;
import org.springframework.web.socket.config.annotation.WebSocketConfigurer;
import org.springframework.web.socket.config.annotation.WebSocketHandlerRegistry;
import org.springframework.web.socket.handler.TextWebSocketHandler;
import org.springframework.web.socket.server.HandshakeInterceptor;
import org.springframework.web.util.UriComponentsBuilder;

import javax.annotation.Resource;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 原生WebSocket配置类
 */
@Slf4j
@Configuration
@EnableWebSocket
public class NativeWebSocketConfig implements WebSocketConfigurer {

    @Resource
    private JwtUtil jwtUtil;
    
    @Resource
    private ConsultService consultService;
    
    @Resource
    private ObjectMapper objectMapper;
    
    // 会话存储，按会话ID分组存储会话
    private final Map<String, Map<String, WebSocketSession>> sessionsByConversation = new ConcurrentHashMap<>();

    @Override
    public void registerWebSocketHandlers(WebSocketHandlerRegistry registry) {
        log.info("注册原生WebSocket处理器");
        
        // 注册一个处理原生WebSocket连接的处理器
        registry.addHandler(chatWebSocketHandler(), "/native-ws")
                .setAllowedOrigins("*")
                .addInterceptors(new WebSocketHandshakeInterceptor());
                
        log.info("原生WebSocket处理器注册完成");
    }
    
    @Bean
    public WebSocketHandler chatWebSocketHandler() {
        return new ChatWebSocketHandler();
    }
    
    /**
     * 自定义原生WebSocket处理器
     */
    public class ChatWebSocketHandler extends TextWebSocketHandler {
        
        @Override
        public void afterConnectionEstablished(WebSocketSession session) throws Exception {
            log.info("原生WebSocket连接已建立，会话ID: {}", session.getId());
            
            // 获取用户ID
            Long userId = (Long) session.getAttributes().get("userId");
            Integer userType = (Integer) session.getAttributes().get("userType");
            String conversationId = (String) session.getAttributes().get("conversationId");
            
            log.info("用户已连接: 用户ID={}, 用户类型={}, 会话ID={}", userId, userType, conversationId);
            
            // 将会话存储到会话管理器中
            if (conversationId != null) {
                Map<String, WebSocketSession> sessions = sessionsByConversation.computeIfAbsent(
                    conversationId, k -> new ConcurrentHashMap<>());
                sessions.put(session.getId(), session);
                log.info("会话已添加到会话管理器, 当前会话组: {}, 会话数: {}", conversationId, sessions.size());
            }
        }
        
        @Override
        protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
            String payload = message.getPayload();
            log.info("收到WebSocket原生消息: {}", payload);
            
            try {
                // 解析消息内容
                MessageDTO messageDTO = objectMapper.readValue(payload, MessageDTO.class);
                
                // 获取会话ID
                String conversationIdStr = (String) session.getAttributes().get("conversationId");
                if (conversationIdStr == null) {
                    log.error("会话ID为空，无法处理消息");
                    sendErrorMessage(session, "会话ID为空");
                    return;
                }
                
                Long conversationId = Long.parseLong(conversationIdStr);
                
                // 保存消息
                Message savedMessage = consultService.saveMessage(conversationId, messageDTO);
                
                if (savedMessage != null) {
                    log.info("原生WebSocket消息保存成功: 消息ID={}", savedMessage.getId());
                    
                    // 广播消息到所有关联的会话
                    broadcastMessageToConversation(conversationIdStr, savedMessage);
                    
                    log.info("消息已广播到所有相关会话");
                } else {
                    log.error("消息保存失败: 会话ID={}", conversationId);
                    sendErrorMessage(session, "消息保存失败");
                }
            } catch (Exception e) {
                log.error("处理WebSocket消息异常", e);
                sendErrorMessage(session, "消息处理出错: " + e.getMessage());
            }
        }
        
        /**
         * 广播消息到同一会话的所有连接
         */
        private void broadcastMessageToConversation(String conversationId, Message message) {
            Map<String, WebSocketSession> sessions = sessionsByConversation.get(conversationId);
            if (sessions != null && !sessions.isEmpty()) {
                try {
                    String messageJson = objectMapper.writeValueAsString(message);
                    TextMessage textMessage = new TextMessage(messageJson);
                    
                    for (WebSocketSession session : sessions.values()) {
                        if (session.isOpen()) {
                            try {
                                session.sendMessage(textMessage);
                                log.debug("消息已发送到会话: {}", session.getId());
                            } catch (IOException e) {
                                log.error("向会话发送消息失败: {}", session.getId(), e);
                            }
                        }
                    }
                } catch (Exception e) {
                    log.error("广播消息序列化失败", e);
                }
            } else {
                log.warn("没有找到会话ID对应的WebSocket连接: {}", conversationId);
            }
        }
        
        private void sendErrorMessage(WebSocketSession session, String errorMessage) {
            try {
                Map<String, String> error = new HashMap<>();
                error.put("error", errorMessage);
                String errorJson = objectMapper.writeValueAsString(error);
                session.sendMessage(new TextMessage(errorJson));
            } catch (IOException e) {
                log.error("发送错误消息失败", e);
            }
        }
        
        @Override
        public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
            log.error("WebSocket传输错误: {}", exception.getMessage(), exception);
        }
        
        @Override
        public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
            log.info("WebSocket连接已关闭: {}, 原因: {}", session.getId(), status);
            
            // 从会话管理器中移除会话
            String conversationId = (String) session.getAttributes().get("conversationId");
            if (conversationId != null) {
                Map<String, WebSocketSession> sessions = sessionsByConversation.get(conversationId);
                if (sessions != null) {
                    sessions.remove(session.getId());
                    log.info("会话已从会话管理器中移除: {}, 剩余会话数: {}", session.getId(), sessions.size());
                    
                    // 如果没有会话了，移除会话组
                    if (sessions.isEmpty()) {
                        sessionsByConversation.remove(conversationId);
                        log.info("会话组已移除: {}", conversationId);
                    }
                }
            }
        }
    }
    
    /**
     * WebSocket握手拦截器
     * 用于在握手前处理token认证
     */
    public class WebSocketHandshakeInterceptor implements HandshakeInterceptor {
        
        @Override
        public boolean beforeHandshake(ServerHttpRequest request, ServerHttpResponse response, 
                                    WebSocketHandler wsHandler, Map<String, Object> attributes) throws Exception {
            log.info("WebSocket握手请求: {}", request.getURI());
            
            try {
                // 从查询参数中获取token
                String query = request.getURI().getQuery();
                Map<String, String> queryParams = UriComponentsBuilder.newInstance()
                        .query(query)
                        .build()
                        .getQueryParams()
                        .toSingleValueMap();
                
                String token = queryParams.get("token");
                log.info("WebSocket Token: {}", token != null ? (token.length() > 10 ? token.substring(0, 10) + "..." : token) : "null");
                
                if (token == null || token.isEmpty()) {
                    log.warn("WebSocket握手失败: 未提供token");
                    return false;
                }
                
                // 验证token
                Claims claims = jwtUtil.getClaimsByToken(token);
                if (claims == null) {
                    log.warn("WebSocket握手失败: token无效");
                    return false;
                }
                
                Long userId = claims.get("userId", Long.class);
                Integer userType = claims.get("userType", Integer.class);
                
                if (userId == null) {
                    log.warn("Token中未包含userId");
                    return false;
                }
                
                log.info("WebSocket连接认证成功，用户ID: {}, 用户类型: {}", userId, userType);
                attributes.put("userId", userId);
                attributes.put("userType", userType);
                
                // 获取会话ID
                String consultId = queryParams.get("consultId");
                if (consultId != null) {
                    attributes.put("conversationId", consultId);
                    log.info("WebSocket连接咨询ID: {}", consultId);
                } else {
                    log.warn("WebSocket连接缺少consultId参数");
                }
                
                return true;
            } catch (Exception e) {
                log.error("WebSocket握手处理异常", e);
                return false;
            }
        }
        
        @Override
        public void afterHandshake(ServerHttpRequest request, ServerHttpResponse response, 
                                 WebSocketHandler wsHandler, Exception exception) {
            // 握手后的处理
            if (exception != null) {
                log.error("WebSocket握手异常", exception);
            } else {
                String uri = request.getURI().toString();
                log.info("WebSocket握手成功, URI: {}", uri);
            }
        }
    }
} 